// WindowCaptureUtil.java
package org.zjx.util;

import com.sun.jna.Memory;
import com.sun.jna.platform.win32.*;
import lombok.extern.slf4j.Slf4j;

import java.awt.image.BufferedImage;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 高效的窗口截图工具类 - 改进资源管理和性能
 */
@Slf4j
public class WindowCaptureUtil {
    // 单例实例
    private static volatile WindowCaptureUtil instance;

    // 窗口上下文缓存
    private final Map<WinDef.HWND, WindowContext> contextMap = new ConcurrentHashMap<>();

    // 清理线程
    private final ScheduledExecutorService cleanupExecutor;
    private final AtomicInteger captureCounter = new AtomicInteger(0);

    // 配置参数
    private static final long CONTEXT_CLEANUP_INTERVAL = 30; // 30分钟清理一次
    private static final long CONTEXT_MAX_IDLE_TIME = 10 * 60 * 1000; // 10分钟空闲时间

    // 私有构造函数
    private WindowCaptureUtil() {
        this.cleanupExecutor = Executors.newSingleThreadScheduledExecutor(r -> {
            Thread t = new Thread(r, "WindowCapture-Cleanup");
            t.setDaemon(true);
            return t;
        });

        // 启动定期清理任务
        startCleanupTask();

        // 注册关闭钩子
        Runtime.getRuntime().addShutdownHook(new Thread(this::cleanupAll, "WindowCapture-Shutdown"));

        log.info("WindowCaptureUtil 初始化完成");
    }

    /**
     * 获取单例实例
     */
    public static synchronized WindowCaptureUtil getInstance() {
        if (instance == null) {
            instance = new WindowCaptureUtil();
        }
        return instance;
    }

    /**
     * 启动清理任务
     */
    private void startCleanupTask() {
        cleanupExecutor.scheduleAtFixedRate(() -> {
            try {
                cleanupStaleContexts();
            } catch (Exception e) {
                log.error("清理窗口上下文失败", e);
            }
        }, CONTEXT_CLEANUP_INTERVAL, CONTEXT_CLEANUP_INTERVAL, TimeUnit.MINUTES);
    }

    private void monitorResourceUsage() {
        Runtime runtime = Runtime.getRuntime();
        long usedMemory = runtime.totalMemory() - runtime.freeMemory();
        if (usedMemory > runtime.maxMemory() * 0.8) {
            log.warn("内存使用率过高，强制清理资源");
            cleanupStaleContexts();
        }
    }

    /**
     * 捕获指定窗口
     */
    public synchronized BufferedImage captureWindow(WinDef.HWND hWnd) {
        if (hWnd == null) {
            throw new IllegalArgumentException("窗口句柄不能为空");
        }
        monitorResourceUsage();

        if (!HwndUtil.haveWindow(hWnd)) {
            log.warn("窗口已失效，无法截图: {}", hWnd);
            removeContext(hWnd);
            return null;
        }

        // 在截图前触发渲染
        triggerWindowRender(hWnd);

        WindowContext context = getOrCreateContext(hWnd);
        int captureId = captureCounter.incrementAndGet();

        try {
            // 检查窗口尺寸是否变化
            if (context.hasWindowSizeChanged()) {
                log.debug("窗口尺寸变化，重新创建资源: {}", hWnd);
                context.recreateResources();
            }

            // 更新最后使用时间
            context.updateLastUsedTime();

            // 执行截图
            BufferedImage image = captureWithPrintWindow(context);

            log.debug("截图成功: ID={}, 窗口={}, 尺寸={}x{}",
                    captureId, hWnd, image.getWidth(), image.getHeight());

            return image;

        } catch (Exception e) {
            log.error("窗口截图失败: ID={}, 窗口={}", captureId, hWnd, e);
            // 发生错误时清理该窗口的上下文
            removeContext(hWnd);
            throw new RuntimeException("窗口截图失败: " + e.getMessage(), e);
        }
    }

    /**
     * 触发窗口渲染（不产生实际点击）
     */
    public void triggerWindowRender(WinDef.HWND hWnd) {
        if (hWnd == null) return;

        try {
            // 发送 WM_ACTIVATE 消息，让游戏退出低功耗模式
            User32.INSTANCE.SendMessage(hWnd, 0x0006, new WinDef.WPARAM(1), null);

            // 短暂等待渲染完成
            Thread.sleep(30);

        } catch (Exception e) {
            log.warn("触发渲染失败: {}", hWnd, e);
        }
    }

    /**
     * 使用PrintWindow方式截图
     */
    private BufferedImage captureWithPrintWindow(WindowContext context) {
        boolean success = User32.INSTANCE.PrintWindow(context.hWnd, context.hdcMemDC, 0x00000002);
        if (!success) {
            throw new RuntimeException("PrintWindow调用失败");
        }

        return convertBitmapToImage(context);
    }

    /**
     * 获取或创建窗口上下文
     */
    private WindowContext getOrCreateContext(WinDef.HWND hWnd){
        return contextMap.computeIfAbsent(hWnd, k -> {
            try {
                WindowContext context = new WindowContext(hWnd);
                context.initialize();
                log.debug("创建窗口上下文: {}", hWnd);
                return context;
            } catch (Exception e) {
                log.error("创建窗口上下文失败: {}", hWnd, e);
                throw new RuntimeException("创建窗口上下文失败", e);
            }
        });
    }

    /**
     * 移除窗口上下文
     */
    private void removeContext(WinDef.HWND hWnd) {
        WindowContext context = contextMap.remove(hWnd);
        if (context != null) {
            context.cleanup();
            log.debug("移除窗口上下文: {}", hWnd);
        }
    }

    /**
     * 转换位图为BufferedImage
     */
    private BufferedImage convertBitmapToImage(WindowContext context){
        // 准备BITMAPINFO结构
        WinGDI.BITMAPINFO bmi = new WinGDI.BITMAPINFO();
        bmi.bmiHeader.biWidth = context.width;
        bmi.bmiHeader.biHeight = -context.height; // 负值表示自上而下
        bmi.bmiHeader.biPlanes = 1;
        bmi.bmiHeader.biBitCount = 32;
        bmi.bmiHeader.biCompression = WinGDI.BI_RGB;
        bmi.bmiHeader.biSizeImage = 0;
        bmi.bmiHeader.biSize = bmi.bmiHeader.size();

        // 计算缓冲区大小
        int bufferSize = context.height * context.width * 4;
        Memory bufferMemory = new Memory(bufferSize);

        // 调用GetDIBits获取像素数据
        int result = GDI32.INSTANCE.GetDIBits(
                context.hdcMemDC,
                context.hBitmap,
                0,
                context.height,
                bufferMemory,
                bmi,
                WinGDI.DIB_RGB_COLORS
        );

        if (result == 0) {
            throw new RuntimeException("GetDIBits调用失败");
        }

        // 读取数据到byte数组
        byte[] buffer = new byte[bufferSize];
        bufferMemory.read(0, buffer, 0, bufferSize);

        // 创建BufferedImage并填充数据
        BufferedImage image = new BufferedImage(context.width, context.height, BufferedImage.TYPE_INT_ARGB);
        int[] pixelBuffer = new int[context.width * context.height];

        // 转换BGRA到ARGB
        for (int y = 0, pixel = 0; y < context.height; y++) {
            for (int x = 0; x < context.width; x++, pixel++) {
                int idx = (y * context.width + x) * 4;
                int blue = buffer[idx] & 0xFF;
                int green = buffer[idx + 1] & 0xFF;
                int red = buffer[idx + 2] & 0xFF;
                int alpha = buffer[idx + 3] & 0xFF;

                int argb = (alpha << 24) | (red << 16) | (green << 8) | blue;
                pixelBuffer[pixel] = argb;
            }
        }

        image.setRGB(0, 0, context.width, context.height, pixelBuffer, 0, context.width);
        return image;
    }

    /**
     * 清理过期的上下文
     */
    private void cleanupStaleContexts() {
        long currentTime = System.currentTimeMillis();
        int removedCount = 0;

        for (Map.Entry<WinDef.HWND, WindowContext> entry : contextMap.entrySet()) {
            WinDef.HWND hWnd = entry.getKey();
            WindowContext context = entry.getValue();

            // 检查窗口是否仍然有效
            if (!HwndUtil.haveWindow(hWnd)) {
                removeContext(hWnd);
                removedCount++;
                continue;
            }

            // 检查上下文是否空闲时间过长
            if (currentTime - context.lastUsedTime > CONTEXT_MAX_IDLE_TIME) {
                removeContext(hWnd);
                removedCount++;
            }
        }

        if (removedCount > 0) {
            log.info("清理了 {} 个过期的窗口上下文", removedCount);
        }
    }

    /**
     * 清理所有资源
     */
    public synchronized void cleanupAll() {
        log.info("开始清理所有窗口截图资源");

        // 停止清理线程
        if (cleanupExecutor != null && !cleanupExecutor.isShutdown()) {
            cleanupExecutor.shutdown();
            try {
                if (!cleanupExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                    cleanupExecutor.shutdownNow();
                }
            } catch (InterruptedException e) {
                cleanupExecutor.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }

        // 清理所有上下文
        int totalContexts = contextMap.size();
        for (WinDef.HWND hWnd : contextMap.keySet()) {
            WindowContext context = contextMap.get(hWnd);
            if (context != null) {
                context.cleanup();
            }
        }
        contextMap.clear();

        log.info("所有窗口截图资源已清理，共清理 {} 个上下文", totalContexts);
    }

    /**
     * 清理指定窗口的资源
     */
    public synchronized void cleanupWindow(WinDef.HWND hWnd) {
        removeContext(hWnd);
        log.debug("窗口资源已清理: {}", hWnd);
    }

    /**
     * 获取统计信息
     */
    public Map<String, Object> getStatistics() {
        Map<String, Object> stats = new java.util.HashMap<>();
        stats.put("contextCount", contextMap.size());
        stats.put("totalCaptures", captureCounter.get());
        stats.put("cleanupExecutorRunning", cleanupExecutor != null && !cleanupExecutor.isShutdown());
        return stats;
    }

    /**
     * 窗口上下文类
     */
    private static class WindowContext {
        final WinDef.HWND hWnd;
        WinDef.HDC hdcWindow;
        WinDef.HDC hdcMemDC;
        WinDef.HBITMAP hBitmap;
        WinNT.HANDLE hOld;
        int width;
        int height;
        long lastUsedTime;

        WindowContext(WinDef.HWND hWnd) {
            this.hWnd = hWnd;
            this.lastUsedTime = System.currentTimeMillis();
        }

        /**
         * 初始化资源
         */
        void initialize() throws Exception {
            // 获取窗口尺寸
            WinDef.RECT rect = new WinDef.RECT();
            if (!User32.INSTANCE.GetWindowRect(hWnd, rect)) {
                throw new RuntimeException("获取窗口尺寸失败");
            }

            width = rect.right - rect.left;
            height = rect.bottom - rect.top;

            if (width <= 0 || height <= 0) {
                throw new RuntimeException("窗口尺寸无效: " + width + "x" + height);
            }

            // 创建设备上下文
            hdcWindow = User32.INSTANCE.GetDC(hWnd);
            if (hdcWindow == null) {
                throw new RuntimeException("获取窗口DC失败");
            }

            // 创建内存DC
            hdcMemDC = GDI32.INSTANCE.CreateCompatibleDC(hdcWindow);
            if (hdcMemDC == null) {
                User32.INSTANCE.ReleaseDC(hWnd, hdcWindow);
                throw new RuntimeException("创建内存DC失败");
            }

            // 创建兼容位图
            hBitmap = GDI32.INSTANCE.CreateCompatibleBitmap(hdcWindow, width, height);
            if (hBitmap == null) {
                cleanup();
                throw new RuntimeException("创建位图失败");
            }

            // 选入内存DC
            hOld = GDI32.INSTANCE.SelectObject(hdcMemDC, hBitmap);
            if (hOld == null) {
                cleanup();
                throw new RuntimeException("选入位图失败");
            }
        }

        /**
         * 检查窗口尺寸是否变化
         */
        boolean hasWindowSizeChanged() {
            WinDef.RECT rect = new WinDef.RECT();
            if (!User32.INSTANCE.GetWindowRect(hWnd, rect)) {
                return false;
            }

            int newWidth = rect.right - rect.left;
            int newHeight = rect.bottom - rect.top;

            return newWidth != width || newHeight != height;
        }

        /**
         * 重新创建资源
         */
        void recreateResources() throws Exception {
            cleanup();
            initialize();
        }

        /**
         * 更新最后使用时间
         */
        void updateLastUsedTime() {
            this.lastUsedTime = System.currentTimeMillis();
        }

        /**
         * 清理资源
         */
        void cleanup() {
            try {
                if (hOld != null && hdcMemDC != null) {
                    GDI32.INSTANCE.SelectObject(hdcMemDC, hOld);
                    hOld = null;
                }

                if (hBitmap != null) {
                    GDI32.INSTANCE.DeleteObject(hBitmap);
                    hBitmap = null;
                }

                if (hdcMemDC != null) {
                    GDI32.INSTANCE.DeleteDC(hdcMemDC);
                    hdcMemDC = null;
                }

                if (hdcWindow != null) {
                    User32.INSTANCE.ReleaseDC(hWnd, hdcWindow);
                    hdcWindow = null;
                }

                width = 0;
                height = 0;

            } catch (Exception e) {
                log.error("清理资源时发生错误: {}", hWnd, e);
            }
        }
    }

    @Override
    protected void finalize() throws Throwable {
        try {
            if (!contextMap.isEmpty()) {
                log.warn("WindowCaptureUtil 在finalize时仍有未清理的资源");
                cleanupAll();
            }
        } finally {
            super.finalize();
        }
    }
}